import torch


class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim1, hidden_dim2, action_dim, action_bound):
        super(PolicyNet, self).__init__()
        self.layer = torch.nn.Sequential(
            torch.nn.Linear(state_dim, hidden_dim1),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim1, hidden_dim2),
            # torch.nn.ReLU(),
            # torch.nn.Linear(hidden_dim2, hidden_dim2),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim2, action_dim),
            # TODO: 如何选择最后一层输出层，从而使得网络的结果能够在正数范围内
            # torch.nn.Sigmoid(),
        )
        # self.action_bound = action_bound  # action_bound是环境可以接受的动作最大值

    def forward(self, x):
        # TODO：如果上面进行更改，这里是不是也得有什么改变
        # return self.layer(x)* self.action_bound
        return self.layer(x)


class QValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim1, hidden_dim2,action_dim):
        super(QValueNet, self).__init__()
        self.layer = torch.nn.Sequential(
            torch.nn.Linear(state_dim+action_dim, hidden_dim1),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim1, hidden_dim2),
            # torch.nn.ReLU(),
            # torch.nn.Linear(hidden_dim2, hidden_dim2),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden_dim2, 1),
        )

    def forward(self, state, a):
        cat = torch.cat([state, a], dim=1)
        return self.layer(cat)
